from langchain.agents import Tool
import requests
import json
from .content_creator_agent import content_creator
from .agent_selector import get_access_token

def content_creator_with_text(user_input):

    generated_text = content_creator(user_input)

    # 使用生成的文本内容直接作为大模型的输入
    system_text = (
        "You are a prompt generator for creating high-quality visuals aimed at digital marketing and content enhancement. Your task is to convert input text into vivid, detailed descriptions ideal for generating images that align precisely with the content’s theme. When the content involves products or services, generate visuals that showcase their unique features in a clear, professional style. Focus on elements such as branding, relevant icons, key product benefits, infographics, or lifestyle scenes to make the content visually appealing and informative. Avoid using any anime, cartoon, or overly stylized effects to ensure the visuals are professional and directly relevant to the content. Please provide your response in English within a 60-word limit.画人的时候，尽量避免画出他们的手部"
    )
    full_input = system_text + generated_text

    url = f"https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro?access_token={get_access_token()}"
    payload = json.dumps({
        "messages": [{"role": "user", "content": full_input}],
        "disable_search": False,
        "system": system_text
    })
    headers = {'Content-Type': 'application/json'}
    response = requests.post(url, headers=headers, data=payload)

    # 输出生成的绘图提示词
    prompt_result = response.json().get("result", "无返回结果")
    return prompt_result

# 定义 content_draw_tool，将 content_creator_with_text 封装为 Tool
content_draw_tool = Tool(
    name="Drawing Prompt Generator",
    func=content_creator_with_text,
    description="直接根据生成的内容生成绘图提示词。"
)

# 定义 generate_image 函数，用于基于提示词生成图像
def generate_image(prompt):
    access_token = get_access_token()
    url = f"https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/text2image/sd_xl?access_token={access_token}"

    payload = json.dumps({
        "prompt": prompt,
        "negative_prompt": "卡通,动漫,二次元",
        "size": "768x768",
        "n": 1,
        "steps": 40,
        "sampler_index": "DPM2"
    })

    headers = {'Content-Type': 'application/json', 'Accept': 'application/json'}

    # 发送请求并打印调试信息
    response = requests.post(url, headers=headers, data=payload)
    print("图像生成 API 响应状态码:", response.status_code)
    print("图像生成 API 响应内容:", response.text)

    if response.status_code == 200:
        data = response.json()
        if 'data' in data and len(data['data']) > 0:
            return data['data'][0].get('b64_image')
        else:
            print("图像生成响应中无图像数据")
            return None
    else:
        print("图像生成请求失败:", response.text)
        return None
